﻿using Microsoft.AspNetCore.Connections;
using System.Collections.Concurrent;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;

namespace ProxyServer.Kestrel
{
    public class CertificateSelector : IServerCertificateSelector
    {
        private readonly ConcurrentDictionary<string, X509Certificate2> _certificates = new ConcurrentDictionary<string, X509Certificate2>();
        private readonly ILogger<CertificateSelector> _logger;

        public CertificateSelector(ILogger<CertificateSelector> logger)
        {
            _logger = logger;
        }

        public X509Certificate2? Select(ConnectionContext? context, string? domainName)
        {
            if (string.IsNullOrWhiteSpace(domainName) || !TryGet(domainName, out var retCert))
            {
                return null;
            }


            _logger.LogDebug("Selected certificate {Thumbprint} for {Domain}", retCert?.Thumbprint, domainName);

            return retCert;
        }

        public bool TryGet(string domainName, out X509Certificate2? certificate)
        {
            if (_certificates.TryGetValue(domainName, out certificate)) return true;
            var wildcardDomainName = _certificates.Keys.FirstOrDefault(n => n.StartsWith("*") && domainName.EndsWith(n[1..]));
            if (wildcardDomainName != null && _certificates.TryGetValue(wildcardDomainName, out certificate))
            {
                return true;
            }
            return false;
        }

        public void Reset(string domain)
        {
            _certificates.TryRemove(domain, out var _);
        }

        public void Add(X509Certificate2 certificate)
        {
            var preloaded = false;

            foreach (var dnsName in X509CertificateHelpers.GetAllDnsNames(certificate))
            {
                var selectedCertificate = AddWithDomainName(_certificates, dnsName, certificate);
                if (!preloaded && selectedCertificate == certificate)
                {
                    preloaded = true;
                    PreloadIntermediateCertificates(certificate);
                }
            }
        }
        
        private void PreloadIntermediateCertificates(X509Certificate2 certificate)
        {
            if (certificate.IsSelfSigned())
            {
                return;
            }

            // workaround for https://github.com/dotnet/aspnetcore/issues/21183
            using var chain = new X509Chain
            {
                ChainPolicy =
                {
                    RevocationMode = X509RevocationMode.NoCheck
                }
            };

            var commonName = X509CertificateHelpers.GetCommonName(certificate);
            try
            {
                var isValid = chain.Build(certificate);

                if (isValid)
                {
                    _logger.LogTrace("Successfully tested certificate chain for {commonName}", commonName);
                    return;
                }

                foreach (var status in chain.ChainStatus)
                {
                    _logger.LogDebug(
                        "Certificate chain error for {commonName}: {Status} - {StatusInformation}",
                        commonName, status.Status, status.StatusInformation);
                }
            }
            catch (CryptographicException ex)
            {
                _logger.LogDebug(ex, "Failed to validate certificate chain for {commonName}", commonName);
            }

            _logger.LogWarning(
                "Failed to validate certificate for {commonName} ({thumbprint}). This could cause an outage of your app.",
                commonName, certificate.Thumbprint);
        }

        /// <summary>
        /// Registers the certificate for usage with domain unless there is already a newer cert for this domain.
        /// </summary>
        /// <param name="certs"></param>
        /// <param name="domainName"></param>
        /// <param name="certificate"></param>
        /// <returns>The certificate current selected to be used for this domain</returns>
        private X509Certificate2 AddWithDomainName(ConcurrentDictionary<string, X509Certificate2> certs, string domainName,
            X509Certificate2 certificate)
        {
            _logger.LogDebug("Adding certificate {Thumbprint} for {domain}", certificate.Thumbprint, domainName);

            return certs.AddOrUpdate(
                domainName,
                certificate,
                (_, currentCert) =>
                {
                    if (currentCert == null || certificate.NotAfter >= currentCert.NotAfter)
                    {
                        _logger.LogDebug("Update certificate {Thumbprint} for {Domain}", certificate.Thumbprint, domainName);
                        return certificate;
                    }

                    _logger.LogDebug("Keep certificate {Thumbprint} for {Domain}", currentCert.Thumbprint, domainName);
                    return currentCert;
                });
        }
    }
}
